// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/map_union_sum_aggregate.h>
#include <pollux/exec/addressable_non_null_value_list.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/strings.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<typename K, typename S>
        struct Accumulator {
            using ValuesMap = typename util::floating_point::HashMapNaNAwareTypeTraits<
                K,
                S,
                AlignedStlAllocator<std::pair<const K, S>, 16> >::Type;
            ValuesMap sums;

            explicit Accumulator(const TypePtr & /*type*/, HashStringAllocator *allocator)
                : sums{AlignedStlAllocator<std::pair<const K, S>, 16>(allocator)} {
            }

            size_t size() const {
                return sums.size();
            }

            void addValues(
                const MapVector *map_vector,
                const VectorPtr &mapKeys,
                const VectorPtr &mapValues,
                vector_size_t row,
                HashStringAllocator *allocator) {
                auto keys = mapKeys->template as<SimpleVector<K> >();
                auto values = mapValues->template as<SimpleVector<S> >();
                auto offset = map_vector->offsetAt(row);
                auto size = map_vector->sizeAt(row);

                for (auto i = 0; i < size; ++i) {
                    // Ignore null map keys.
                    if (!keys->is_null_at(offset + i)) {
                        auto key = keys->value_at(offset + i);
                        addValue(key, values, offset + i, values->type_kind());
                    }
                }
            }

            void addValue(
                K key,
                const SimpleVector<S> *mapValues,
                vector_size_t row,
                TypeKind valueKind) {
                if (mapValues->is_null_at(row)) {
                    sums[key] += 0;
                } else {
                    auto value = mapValues->value_at(row);

                    if constexpr (std::is_same_v<S, double> || std::is_same_v<S, float>) {
                        sums[key] += value;
                    } else {
                        S checkedSum;
                        auto overflow = __builtin_add_overflow(sums[key], value, &checkedSum);

                        if (UNLIKELY(overflow)) {
                            auto errorValue = (int128_t(sums[key]) + int128_t(value));

                            if (errorValue < 0) {
                                POLLUX_ARITHMETIC_ERROR(
                                    "Value {} is less than {}",
                                    errorValue,
                                    std::numeric_limits<S>::min());
                            } else {
                                POLLUX_ARITHMETIC_ERROR(
                                    "Value {} exceeds {}",
                                    errorValue,
                                    std::numeric_limits<S>::max());
                            }
                        }
                        sums[key] = checkedSum;
                    }
                }
            }

            vector_size_t extractValues(
                VectorPtr &mapKeys,
                VectorPtr &mapValues,
                vector_size_t offset) {
                auto keys = mapKeys->as_flat_vector<K>();
                auto values = mapValues->as_flat_vector<S>();

                auto index = offset;
                for (const auto &[key, sum]: sums) {
                    keys->set(index, key);
                    values->set(index, sum);

                    ++index;
                }

                return sums.size();
            }
        };

        template<typename S>
        struct StringViewAccumulator {
            Accumulator<StringView, S> base;

            Strings strings;

            explicit StringViewAccumulator(
                const TypePtr &type,
                HashStringAllocator *allocator)
                : base{type, allocator} {
            }

            size_t size() const {
                return base.size();
            }

            void addValues(
                const MapVector *map_vector,
                const VectorPtr &mapKeys,
                const VectorPtr &mapValues,
                vector_size_t row,
                HashStringAllocator *allocator) {
                auto keys = mapKeys->template as<SimpleVector<StringView> >();
                auto values = mapValues->template as<SimpleVector<S> >();
                auto offset = map_vector->offsetAt(row);
                auto size = map_vector->sizeAt(row);

                for (auto i = 0; i < size; ++i) {
                    // Ignore null map keys.
                    if (!keys->is_null_at(offset + i)) {
                        auto key = keys->value_at(offset + i);

                        if (!key.isInline()) {
                            auto it = base.sums.find(key);
                            if (it != base.sums.end()) {
                                key = it->first;
                            } else {
                                key = strings.append(key, *allocator);
                            }
                        }

                        base.addValue(key, values, offset + i, values->type_kind());
                    }
                }
            }

            vector_size_t extractValues(
                VectorPtr &mapKeys,
                VectorPtr &mapValues,
                vector_size_t offset) {
                return base.extractValues(mapKeys, mapValues, offset);
            }
        };

        /// Maintains a map with keys of type array, map or struct.
        template<typename V>
        struct ComplexTypeAccumulator {
            using ValueMap = melon::F14FastMap<
                AddressableNonNullValueList::Entry,
                V,
                AddressableNonNullValueList::Hash,
                AddressableNonNullValueList::EqualTo,
                AlignedStlAllocator<
                    std::pair<const AddressableNonNullValueList::Entry, V>,
                    16> >;

            /// A set of pointers to values stored in AddressableNonNullValueList.
            ValueMap sums;

            /// Stores unique non-null keys.
            AddressableNonNullValueList serializedKeys;

            ComplexTypeAccumulator(const TypePtr &type, HashStringAllocator *allocator)
                : sums{
                    0,
                    AddressableNonNullValueList::Hash{},
                    AddressableNonNullValueList::EqualTo{type->as_map().keyType()},
                    AlignedStlAllocator<
                        std::pair<const AddressableNonNullValueList::Entry, V>,
                        16>(allocator)
                } {
            }

            void addValues(
                const MapVector *map_vector,
                const VectorPtr &mapKeys,
                const VectorPtr &mapValues,
                vector_size_t row,
                HashStringAllocator *allocator) {
                auto offset = map_vector->offsetAt(row);
                auto size = map_vector->sizeAt(row);
                auto values = mapValues->template as<SimpleVector<V> >();

                for (auto i = 0; i < size; ++i) {
                    if (!mapKeys->is_null_at(offset + i) && (offset + i) < mapKeys->size()) {
                        // Get entry and value to add.
                        auto entry =
                                serializedKeys.append(*mapKeys.get(), offset + i, allocator);
                        auto value =
                                (values->is_null_at(offset + i)) ? 0 : values->value_at(offset + i);

                        // New entry.
                        if (!sums.contains(entry)) {
                            sums[entry] = value;
                        } else {
                            // Existing entry.
                            if constexpr (std::is_same_v<V, double> || std::is_same_v<V, float>) {
                                sums[entry] += value;
                            } else {
                                V checkedSum;
                                auto overflow =
                                        __builtin_add_overflow(sums[entry], value, &checkedSum);

                                if (UNLIKELY(overflow)) {
                                    auto errorValue = (int128_t(sums[entry]) + int128_t(value));

                                    if (errorValue < 0) {
                                        POLLUX_ARITHMETIC_ERROR(
                                            "Value {} is less than {}",
                                            errorValue,
                                            std::numeric_limits<V>::min());
                                    } else {
                                        POLLUX_ARITHMETIC_ERROR(
                                            "Value {} exceeds {}",
                                            errorValue,
                                            std::numeric_limits<V>::max());
                                    }
                                }
                                sums[entry] = checkedSum;
                                serializedKeys.removeLast(entry);
                            }
                        }
                    }
                }
            }

            vector_size_t extractValues(
                VectorPtr &mapKeys,
                VectorPtr &mapValues,
                vector_size_t offset) {
                auto values = mapValues->as_flat_vector<V>();
                auto index = offset;

                for (const auto &[position, count]: sums) {
                    AddressableNonNullValueList::read(position, *mapKeys.get(), index);
                    values->set(index, count);
                    ++index;
                }

                return sums.size();
            }

            size_t size() const {
                return sums.size();
            }

            void free(HashStringAllocator &allocator) {
                serializedKeys.free(allocator);
                std::destroy_at(&sums);
            }
        };

        // Defines unique accumulators dependent on type.
        template<typename K, typename S>
        struct AccumulatorTypeTraits {
            using AccumulatorType = Accumulator<K, S>;
        };

        template<typename S>
        struct AccumulatorTypeTraits<StringView, S> {
            using AccumulatorType = StringViewAccumulator<S>;
        };

        template<typename V>
        struct AccumulatorTypeTraits<ComplexType, V> {
            using AccumulatorType = ComplexTypeAccumulator<V>;
        };

        template<typename S>
        struct AccumulatorTypeTraits<UnknownValue, S> {
            using AccumulatorType = Accumulator<int32_t, S>;
        };

        // Defines common aggregator.
        template<typename K, typename S>
        class MapUnionSumAggregate : public exec::Aggregate {
        public:
            explicit MapUnionSumAggregate(TypePtr resultType)
                : Aggregate(std::move(resultType)) {
            }

            using AccumulatorType = typename AccumulatorTypeTraits<K, S>::AccumulatorType;

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(AccumulatorType);
            }

            bool isFixedSize() const override {
                return false;
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto map_vector = (*result)->as<MapVector>();
                POLLUX_CHECK(map_vector);
                map_vector->resize(numGroups);

                auto mapKeysPtr = map_vector->mapKeys();
                auto mapValuesPtr = map_vector->mapValues();

                auto numElements = countElements(groups, numGroups);
                map_vector->mapValues()->as<FlatVector<S> >()->resize(numElements);

                // ComplexType cannot be resized the same.
                if constexpr (!std::is_same_v<K, ComplexType>) {
                    map_vector->mapKeys()->as<FlatVector<K> >()->resize(numElements);
                } else {
                    map_vector->mapKeys()->resize(numElements);
                }

                auto rawNulls = map_vector->mutable_raw_nulls();
                vector_size_t offset = 0;
                for (auto i = 0; i < numGroups; ++i) {
                    char *group = groups[i];
                    if (isNull(group)) {
                        bits::set_null(rawNulls, i, true);
                        map_vector->setOffsetAndSize(i, 0, 0);
                    } else {
                        clearNull(rawNulls, i);

                        auto mapSize = value<AccumulatorType>(group)->extractValues(
                            mapKeysPtr, mapValuesPtr, offset);
                        map_vector->setOffsetAndSize(i, offset, mapSize);
                        offset += mapSize;
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodedMaps_.decode(*args[0], rows);
                auto map_vector = decodedMaps_.base()->template as<MapVector>();
                auto mapKeys = map_vector->mapKeys();
                auto mapValues = map_vector->mapValues();

                rows.applyToSelected([&](auto row) {
                    if (!decodedMaps_.is_null_at(row)) {
                        auto *group = groups[row];
                        clearNull(group);

                        auto tracker = trackRowSize(group);
                        auto groupMap = value<AccumulatorType>(group);
                        addMap(*groupMap, map_vector, mapKeys, mapValues, row);
                    }
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                decodedMaps_.decode(*args[0], rows);
                auto map_vector = decodedMaps_.base()->template as<MapVector>();
                auto mapKeys = map_vector->mapKeys();
                auto mapValues = map_vector->mapValues();

                auto groupMap = value<AccumulatorType>(group);

                auto tracker = trackRowSize(group);
                rows.applyToSelected([&](auto row) {
                    if (!decodedMaps_.is_null_at(row)) {
                        clearNull(group);
                        addMap(*groupMap, map_vector, mapKeys, mapValues, row);
                    }
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                addRawInput(groups, rows, args, false);
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                addSingleGroupRawInput(group, rows, args, false);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                setAllNulls(groups, indices);
                for (auto index: indices) {
                    new(groups[index] + offset_) AccumulatorType{resultType_, allocator_};
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                if constexpr (std::is_same_v<K, StringView>) {
                    for (auto *group: groups) {
                        if (isInitialized(group) && !isNull(group)) {
                            value<AccumulatorType>(group)->strings.free(*allocator_);
                        }
                    }
                } else if constexpr (std::is_same_v<K, ComplexType>) {
                    for (auto *group: groups) {
                        if (isInitialized(group) && !isNull(group)) {
                            value<AccumulatorType>(group)->free(*allocator_);
                        }
                    }
                }
                destroyAccumulators<AccumulatorType>(groups);
            }

        private:
            void addMap(
                AccumulatorType &groupMap,
                const MapVector *map_vector,
                const VectorPtr &mapKeys,
                const VectorPtr &mapValues,
                vector_size_t row) const {
                auto decodedRow = decodedMaps_.index(row);
                groupMap.addValues(map_vector, mapKeys, mapValues, decodedRow, allocator_);
            }

            vector_size_t countElements(char **groups, int32_t numGroups) const {
                vector_size_t size = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    size += value<AccumulatorType>(groups[i])->size();
                }
                return size;
            }

            DecodedVector decodedMaps_;
        };

        template<typename K>
        std::unique_ptr<exec::Aggregate> createMapUnionSumAggregate(
            TypeKind valueKind,
            const TypePtr &resultType) {
            switch (valueKind) {
                case TypeKind::TINYINT:
                    return std::make_unique<MapUnionSumAggregate<K, int8_t> >(resultType);
                case TypeKind::SMALLINT:
                    return std::make_unique<MapUnionSumAggregate<K, int16_t> >(resultType);
                case TypeKind::INTEGER:
                    return std::make_unique<MapUnionSumAggregate<K, int32_t> >(resultType);
                case TypeKind::BIGINT:
                    return std::make_unique<MapUnionSumAggregate<K, int64_t> >(resultType);
                case TypeKind::REAL:
                    return std::make_unique<MapUnionSumAggregate<K, float> >(resultType);
                case TypeKind::DOUBLE:
                    return std::make_unique<MapUnionSumAggregate<K, double> >(resultType);
                default:
                    POLLUX_UNREACHABLE(
                        "Unexpected value type {}", mapTypeKindToName(valueKind));
            }
        }
    } // namespace

    void registerMapUnionSumAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        const std::vector<std::string> valueTypes = {
            "tinyint", "smallint", "integer", "bigint", "double", "real"
        };

        // Add all allowed signatures.
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        for (auto valueType: valueTypes) {
            signatures.push_back(
                exec::AggregateFunctionSignatureBuilder()
                .typeVariable("K")
                .returnType(fmt::format("map(K,{})", valueType))
                .intermediateType(fmt::format("map(K,{})", valueType))
                .argumentType(fmt::format("map(K,{})", valueType))
                .build());
        }

        auto name = prefix + kMapUnionSum;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step /*step*/,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(argTypes.size(), 1);
                POLLUX_CHECK(argTypes[0]->is_map());
                auto &mapType = argTypes[0]->as_map();
                auto keyTypeKind = mapType.keyType()->kind();
                auto valueTypeKind = mapType.valueType()->kind();

                switch (keyTypeKind) {
                    case TypeKind::BOOLEAN:
                        return createMapUnionSumAggregate<bool>(valueTypeKind, resultType);
                    case TypeKind::TINYINT:
                        return createMapUnionSumAggregate<int8_t>(
                            valueTypeKind, resultType);
                    case TypeKind::SMALLINT:
                        return createMapUnionSumAggregate<int16_t>(
                            valueTypeKind, resultType);
                    case TypeKind::INTEGER:
                        return createMapUnionSumAggregate<int32_t>(
                            valueTypeKind, resultType);
                    case TypeKind::BIGINT:
                        return createMapUnionSumAggregate<int64_t>(
                            valueTypeKind, resultType);
                    case TypeKind::REAL:
                        return createMapUnionSumAggregate<float>(valueTypeKind, resultType);
                    case TypeKind::DOUBLE:
                        return createMapUnionSumAggregate<double>(
                            valueTypeKind, resultType);
                    case TypeKind::TIMESTAMP:
                        return createMapUnionSumAggregate<Timestamp>(
                            valueTypeKind, resultType);
                    case TypeKind::VARBINARY:
                    case TypeKind::VARCHAR:
                        return createMapUnionSumAggregate<StringView>(
                            valueTypeKind, resultType);
                    case TypeKind::ARRAY:
                    case TypeKind::MAP:
                    case TypeKind::ROW:
                        return createMapUnionSumAggregate<ComplexType>(
                            valueTypeKind, resultType);
                    case TypeKind::UNKNOWN:
                        return createMapUnionSumAggregate<UnknownValue>(
                            valueTypeKind, resultType);
                    default:
                        if (mapType.keyType()->isDecimal()) {
                            return createMapUnionSumAggregate<int128_t>(
                                valueTypeKind, resultType);
                        }
                        POLLUX_UNREACHABLE(
                            "Unexpected key type {}", mapTypeKindToName(keyTypeKind));
                }
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
